#!/bin/bash
# ==============================================================================
# GSICoreAnalysis.jl - Main Analysis Runner Script
# ==============================================================================
# This script provides a comprehensive interface to run GSI atmospheric data
# assimilation analysis with various options and configurations.

set -e  # Exit on any error

# Color codes for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Script configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LOG_DIR="$SCRIPT_DIR/logs"
RESULT_DIR="$SCRIPT_DIR/results"

# Julia configuration
JULIA_PROJECT="$SCRIPT_DIR"
JULIA_THREADS="${JULIA_NUM_THREADS:-4}"
JULIA_HOME="$HOME/usr/julia"
JULIA_BINARY="$JULIA_HOME/bin/julia"

# Fallback to system Julia if custom installation doesn't exist
if [[ ! -f "$JULIA_BINARY" ]]; then
    JULIA_BINARY="/usr/local/bin/julia"
fi

# Create necessary directories
mkdir -p "$LOG_DIR" "$RESULT_DIR"

# ==============================================================================
# HELPER FUNCTIONS
# ==============================================================================

print_header() {
    echo -e "${BLUE}================================================${NC}"
    echo -e "${BLUE}    GSICoreAnalysis.jl - Main Runner${NC}"
    echo -e "${BLUE}================================================${NC}"
    echo "Script Directory: $SCRIPT_DIR"
    echo "Julia Threads: $JULIA_THREADS"
    echo "Log Directory: $LOG_DIR"
    echo "Results Directory: $RESULT_DIR"
    echo ""
}

print_usage() {
    echo "Usage: $0 [OPTIONS] [ANALYSIS_TYPE]"
    echo ""
    echo "ANALYSIS_TYPES:"
    echo "  quick          - Quick demonstration (5-10 seconds)"
    echo "  basic          - Basic GSI analysis workflow"
    echo "  regional       - Regional GSI test case (default)"
    echo "  comprehensive  - Comprehensive analysis with all features"
    echo "  lorenz         - Lorenz-96 test case"
    echo "  performance    - Performance benchmarking"
    echo "  letkf          - LETKF workflow demonstration"
    echo ""
    echo "OPTIONS:"
    echo "  -h, --help     - Show this help message"
    echo "  -v, --verbose  - Enable verbose output"
    echo "  -t, --threads N - Set number of Julia threads (default: 4)"
    echo "  -l, --log FILE - Custom log file name"
    echo "  --dry-run      - Show commands without executing"
    echo "  --visualize    - Enable visualization (if supported)"
    echo "  --no-deps      - Skip dependency installation check"
    echo ""
    echo "EXAMPLES:"
    echo "  $0                           # Run default regional analysis"
    echo "  $0 quick                     # Quick demo"
    echo "  $0 --visualize regional      # Regional analysis with plots"
    echo "  $0 -t 8 comprehensive        # Comprehensive with 8 threads"
    echo ""
}

check_julia() {
    if [[ ! -f "$JULIA_BINARY" ]]; then
        echo -e "${RED}ERROR: Julia not found at $JULIA_BINARY${NC}" >&2
        echo "Please install Julia to $JULIA_HOME or ensure Julia is in PATH" >&2
        exit 1
    fi

    # Test if Julia actually works (check for library issues)
    local julia_test_output
    local julia_exit_code
    echo "Testing Julia at: $JULIA_BINARY"

    # Use a more robust timeout approach
    if timeout 5s bash -c "\"$JULIA_BINARY\" --version" >/tmp/julia_test.out 2>&1; then
        julia_exit_code=0
        julia_test_output=$(cat /tmp/julia_test.out)
    else
        julia_exit_code=1
        julia_test_output=$(cat /tmp/julia_test.out 2>/dev/null || echo "Julia test failed or timed out")
    fi
    rm -f /tmp/julia_test.out

    if [[ $julia_exit_code -ne 0 ]]; then
        echo -e "${YELLOW}⚠️ Custom Julia has issues: $julia_test_output${NC}"
        if [[ -f "/usr/local/bin/julia" ]]; then
            echo -e "${BLUE}🔄 Falling back to system Julia...${NC}"
            JULIA_BINARY="/usr/local/bin/julia"
            julia_test_output=$("$JULIA_BINARY" --version 2>&1)
            julia_exit_code=$?
            if [[ $julia_exit_code -ne 0 ]]; then
                echo -e "${RED}ERROR: System Julia also has issues: $julia_test_output${NC}" >&2
                exit 1
            fi
        else
            echo -e "${RED}ERROR: No working Julia installation found${NC}" >&2
            exit 1
        fi
    fi

    local julia_version=$(echo "$julia_test_output" | cut -d' ' -f3)
    if [[ "$JULIA_BINARY" == "/usr/local/bin/julia" ]]; then
        echo -e "${GREEN}✓ Julia found: version $julia_version (using system Julia)${NC}"
    else
        echo -e "${GREEN}✓ Julia found: version $julia_version at $JULIA_BINARY${NC}"
    fi
}

check_dependencies() {
    if [[ "$SKIP_DEPS" == "true" ]]; then
        echo -e "${YELLOW}⚠ Skipping dependency check${NC}"
        return 0
    fi

    echo "Checking Julia dependencies..."
    cd "$SCRIPT_DIR"

    # Check if Project.toml exists
    if [[ ! -f "Project.toml" ]]; then
        echo -e "${RED}ERROR: Project.toml not found${NC}" >&2
        exit 1
    fi

    # Install/update dependencies
    "$JULIA_BINARY" --project=. -e "using Pkg; Pkg.instantiate(); Pkg.precompile()" || {
        echo -e "${RED}ERROR: Failed to install dependencies${NC}" >&2
        exit 1
    }

    echo -e "${GREEN}✓ Dependencies checked${NC}"
}

run_analysis() {
    local analysis_type="$1"
    local log_file="$LOG_DIR/gsi_${analysis_type}_$(date +%Y%m%d_%H%M%S).log"

    if [[ -n "$CUSTOM_LOG" ]]; then
        log_file="$LOG_DIR/$CUSTOM_LOG"
    fi

    echo -e "${BLUE}Starting $analysis_type analysis...${NC}"
    echo "Log file: $log_file"
    echo ""

    local cmd
    local script_file

    case "$analysis_type" in
        "quick")
            script_file="demos/quick_demo.jl"
            ;;
        "basic")
            script_file="examples/basic_example.jl"
            ;;
        "regional")
            script_file="demos/gsi_test_case_demo.jl"
            ;;
        "comprehensive")
            script_file="demos/comprehensive_demo.jl"
            ;;
        "lorenz")
            script_file="test/test_lorenz96.jl"
            ;;
        "performance")
            script_file="demos/performance_demo.jl"
            ;;
        "letkf")
            script_file="demos/letkf_workflow_demo.jl"
            ;;
        *)
            echo -e "${RED}ERROR: Unknown analysis type: $analysis_type${NC}" >&2
            exit 1
            ;;
    esac

    if [[ ! -f "$script_file" ]]; then
        echo -e "${RED}ERROR: Script not found: $script_file${NC}" >&2
        exit 1
    fi

    # Build Julia command
    cmd="$JULIA_BINARY --project=. --threads=$JULIA_THREADS"

    if [[ "$VERBOSE" == "true" ]]; then
        cmd="$cmd --verbose"
    fi

    cmd="$cmd $script_file"

    if [[ "$DRY_RUN" == "true" ]]; then
        echo -e "${YELLOW}DRY RUN - Would execute:${NC}"
        echo "$cmd"
        echo "Log: $log_file"
        return 0
    fi

    # Set environment variables
    export JULIA_NUM_THREADS="$JULIA_THREADS"
    export GSI_LOG_FILE="$log_file"

    # Execute the command
    local start_time=$(date +%s)
    echo -e "${GREEN}Executing: $cmd${NC}"

    if [[ "$VERBOSE" == "true" ]]; then
        $cmd 2>&1 | tee "$log_file"
    else
        $cmd > "$log_file" 2>&1
    fi

    local exit_code=$?
    local end_time=$(date +%s)
    local duration=$((end_time - start_time))

    if [[ $exit_code -eq 0 ]]; then
        echo -e "${GREEN}✓ Analysis completed successfully${NC}"
        echo "Duration: ${duration}s"
        echo "Log: $log_file"

        # Check for generated results
        local result_count=0
        if [[ -d "tests/drp4dvar_visualizations" ]]; then
            result_count=$(find tests/drp4dvar_visualizations -name "*.png" -o -name "*.md" | wc -l)
            if [[ $result_count -gt 0 ]]; then
                echo -e "${GREEN}Generated $result_count visualization files${NC}"
                echo "Results: $(pwd)/tests/drp4dvar_visualizations/"
            fi
        fi

    else
        echo -e "${RED}✗ Analysis failed with exit code: $exit_code${NC}" >&2
        echo "Check log file: $log_file" >&2
        exit $exit_code
    fi
}

# ==============================================================================
# MAIN SCRIPT
# ==============================================================================

# Default values
ANALYSIS_TYPE="regional"
VERBOSE="false"
DRY_RUN="false"
VISUALIZE="false"
SKIP_DEPS="false"
CUSTOM_LOG=""

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help)
            print_usage
            exit 0
            ;;
        -v|--verbose)
            VERBOSE="true"
            shift
            ;;
        -t|--threads)
            JULIA_THREADS="$2"
            shift 2
            ;;
        -l|--log)
            CUSTOM_LOG="$2"
            shift 2
            ;;
        --dry-run)
            DRY_RUN="true"
            shift
            ;;
        --visualize)
            VISUALIZE="true"
            shift
            ;;
        --no-deps)
            SKIP_DEPS="true"
            shift
            ;;
        quick|basic|regional|comprehensive|lorenz|performance|letkf)
            ANALYSIS_TYPE="$1"
            shift
            ;;
        *)
            echo -e "${RED}ERROR: Unknown option: $1${NC}" >&2
            print_usage
            exit 1
            ;;
    esac
done

# Main execution
print_header

# Validation and setup
check_julia
check_dependencies

# Handle visualization option
if [[ "$VISUALIZE" == "true" ]]; then
    echo "Installing visualization dependencies..."
    "$JULIA_BINARY" --project=. -e "using Pkg; Pkg.add([\"Plots\", \"StatsPlots\", \"PlotlyJS\", \"StatsBase\"])"
fi

# Run the analysis
echo -e "${BLUE}Running GSI Analysis: $ANALYSIS_TYPE${NC}"
echo "Configuration:"
echo "  Analysis Type: $ANALYSIS_TYPE"
echo "  Julia Threads: $JULIA_THREADS"
echo "  Verbose Mode: $VERBOSE"
echo "  Visualization: $VISUALIZE"
echo ""

run_analysis "$ANALYSIS_TYPE"

echo ""
echo -e "${GREEN}Script completed successfully!${NC}"
echo "Check logs in: $LOG_DIR"
if [[ -d "tests/drp4dvar_visualizations" ]]; then
    echo "Check results in: $(pwd)/tests/drp4dvar_visualizations/"
fi